#-*- coding:utf-8 -*-

# json模块
import json
# 随机数模块
import random
# 时间模块
import time, os, math
# 高性能数组模块
import numpy as np
from vno import *
from vno_data import *
# from scipy.spatial import distance
# 遗传算法模块inspyred
from inspyred import ec
from inspyred import swarm
from inspyred.ec import terminators
from inspyred.ec import observers
from inspyred.ec import analysis
import inspyred
import itertools

# 欧式距离
from numpy.linalg import norm as np_linalg_norm

# 表格美化打印输出模块
from prettytable import PrettyTable

# 适应值曲线绘制模块
import ea_utils

import multiprocessing  # 多进程
import platform

# 小数点个数
DOUBLE_PRECISION = 6

# 清屏
def clear_dis():
    if platform.uname()[0] == 'Windows':
        os.system('cls')
    else:
        os.system('clear')

def get_dict_value(kwargs, key, value):
    if key in kwargs:
        return kwargs[key]
    else:
        return value

def creat_dir(dirs):
    if not os.path.exists(dirs):
        os.makedirs(dirs)

# 选择给定id的风压、风量、节点压能值
def sel_features(sel_fs,origin_fs):
    if len(sel_fs) < 1:
        return origin_fs
    ret = {}
    for k in sel_fs.keys():
        if k not in origin_fs.keys():
            print(sel_fs,origin_fs)
            return None
        ret[k] = origin_fs[k]
    return ret

# 选择给定两点之间的压差，在网络解算计算节点压能时不考虑风机负压时直接用节点压能相减即可
def sel_delta_h(vnet,sensor_delta_h):
    dg = vnet.graph()
    # 压差,格式为{'高压节点id-低压节点id':传感器(或给定)的压差}
    delta_h_dict = {}
    for k,value in sensor_delta_h.items():
        lu_ids = k.split('-')
        high_h = 0.0
        low_h = 0.0
        for vv in dg.vs:
            # 排除虚拟节点
            if (vnet.sn == vv.index)  or  (vnet.tn == vv.index):
                continue
            # 高压节点
            if vv['id'] == lu_ids[0]:
                high_h = vv['p']
            # 低压节点
            if vv['id'] == lu_ids[1]:
                low_h = vv['p']
            if (high_h != 0.0) and (low_h != 0.0):
                break
        # 所有节点压能都小于等于0，所以高压点减去低压点为正数
        delta_h_dict[k] = high_h - low_h
    return delta_h_dict

def sel_struct_h(edges,sel_st_h,edge_ids,delta_rs):
    st_hs = {}
    for st_h_eid,st_h in sel_st_h.items():
        for eid,dr in zip(edge_ids,delta_rs):
            if eid == st_h_eid:
                for e in edges:
                    if e['id'] == st_h_eid:
                        st_hs[st_h_eid] = (e['r']+e['adjust_r']-e['r0']-dr)*e['q']*abs(e['q'])
                        break
    return st_hs

def sel_fan_h(edges,sel_fan_h):
    fan_hs = {}
    for fan_h_evid,st_h in sel_fan_h.items():
        fan_h_eid = fan_h_evid.split('-')[-1]
        for e in edges:
            if e['id'] == fan_h_eid:
                fan_hs[fan_h_evid] = h0(e)
                break
    return fan_hs

'''
vnet,网络对象
edge_ids,分支id
struct_eids,构筑物id 
delta_rs,分支id对应的风阻变化（编码） 
strcut_rs,构筑物id对应的风阻变化（编码） 
sel_params,传感器参数
is_print,是否打印解算结果
all_q,选择所有分支风量，还是传感器对应分支的风量,yes:所有风量
'''
def each_vno(vnet,edge_ids,struct_eids, delta_rs, strcut_rs, sel_params = None, is_print = False,all_q = False):
    dg = vnet.graph()
    # 改变分支风阻
    for eid,dr in zip(edge_ids,delta_rs):
        for e in dg.es:
            if e['id'] == eid:
                e['adjust_r'] = dr
                # print(eid)
                # 风阻变到很小
                if R(e) <= 1e-6:
                    return None,None,None,None
                break
    # 增加构筑物变化风阻
    for eid,dr in zip(struct_eids,strcut_rs):
        for e in dg.es:
            if e['id'] == eid:
                e['adjust_r'] += dr
                # print(eid,e['adjust_r'])
                # 风阻变到很小
                if R(e) <= 1e-6:
                    return None,None,None,None
                break

    ret = vno(vnet)
    # print([e['adjust_r'] for e in dg.es])
    if not ret:
        # 恢复风阻
        for e in dg.es:
            if e['id'] in edge_ids:
                # print(e['id'],e['adjust_r'])
                e['adjust_r'] = 0
        # print('通风网络解算失败!')
        return None,None,None,None
    # 风量传感器数据
    q_dict = {}
    # 风机压差传感器数据,# 格式：{'末节点id-分支id':压差}
    fan_h_dict = {}
    # 构筑物压差传感器数据
    st_h_dict = {}
    # 压差,格式为{'高压节点id-低压节点id':传感器(或给定)的压差}
    delta_h_dict = {}
    for ee in dg.es:
        if is_zero_edge(ee):
            continue
        # print('xx',h0(ee))
        q_dict[ee['id']] = ee['q']

    if not all_q:
        q_dict = sel_features(sel_params['q'],q_dict)
    if 'delta_h' in sel_params:
        delta_h_dict = sel_delta_h(vnet,sel_params['delta_h'])
    if 'st_h' in sel_params:
        st_h_dict = sel_struct_h(dg.es,sel_params['st_h'],edge_ids,delta_rs)
    if 'fan_h' in sel_params:
        fan_h_dict = sel_fan_h(dg.es,sel_params['fan_h'])

    if is_print:
        print_network(vnet, msg='网络解算结果')
        print_mrp_fans(vnet)
        print_mrp(vnet)

    # 恢复风阻
    for e in dg.es:
        if e['id'] in edge_ids:
            e['adjust_r'] = 0
    return q_dict, delta_h_dict, st_h_dict, fan_h_dict

# 计算相对误差
def cacul_rela_erros(tv,pv,weight=1):
    if abs(tv) < 1e-6:
        tv = 0.05
    e = weight*abs(abs(tv-pv)/tv)
    return e

class Problem():
    def __init__(self,args):
        self.lb = args.get('lb',0)
        self.ub = args.get('ub',1)
        self.bounder = ec.Bounder(self.lb, self.ub)
        self.var_types = args['var_types']
        self.var_num = len(self.var_types)
        self.pro_datas = args['pro_datas']
        self.edge_ids = args['edge_ids']
        self.edge_num = len(self.edge_ids)
        self.struct_eids = args['struct_eids']
        self.struct_num = len(self.struct_eids)
        self.vnet = args['vnet']
        self.sel_params = args['sel_params']
        self.Q0 = args['Q0']
        self.H0 = args['H0']
        self.P0 = args['P0']
        self.weight = args['weight']

    # 返回变量的上下限
    def bounder(self):
        return ec.Bounder(self.lb, self.ub)
        
    # 随机生成一个在上下限范围内的编码
    # random -- 随机数生成器对象
    def rand(self, rng):
        lub = [rng.uniform(self.lb[i], self.ub[i]) for i in range(self.var_num)]
        # edge_id = rng.choice(self.edge_ids)
        # lub[0] = edge_id
        return lub
    
    # 将个体编码(取值范围[0,1])解码为"真实变量"(取值范围[lb, ub])
    # 数学计算公式: f=(b-a)*x+a, 其中x范围[0,1], f范围[a, b]
    def decode(self, X):
        # return [self.g(i, x) for i,x in enumerate(X)]
        d_x = []
        for i,x in enumerate(X):
            if self.var_types[i] == 0:
                d_x.append(int(x))
            else:
                d_x.append(round(x,DOUBLE_PRECISION))
        return d_x

    # 相对误差
    def rela_erros(self,X):
        # 分支编码
        X1 = X[0:self.edge_num]
        # 构筑物编码
        X2 = X[self.edge_num:self.struct_num+self.edge_num]
        q_dict,delta_h_dict, st_h_dict, fan_h_dict = each_vno(self.vnet,self.edge_ids,self.struct_eids,X1,X2,sel_params=self.sel_params)
        if (q_dict is None) or (delta_h_dict is None):
            return 1e6       
        sel_q = self.sel_params['q']
        sel_delta_h = self.sel_params['delta_h']
        sel_st_h = self.sel_params['st_h']
        sel_fan_h = self.sel_params['fan_h']
        edge_set = self.pro_datas['edge_set']
        node_set = self.pro_datas['node_set']
        eq = []
        for k in sel_q:
            eq.append(cacul_rela_erros(sel_q[k],q_dict[k]))
        edh = []
        for k in sel_delta_h:
            edh.append(cacul_rela_erros(sel_delta_h[k],delta_h_dict[k]))
        sth = []
        for k in sel_st_h:
            sth.append(cacul_rela_erros(sel_st_h[k],st_h_dict[k]))
        fanh = []
        for k in sel_fan_h:
            fanh.append(cacul_rela_erros(sel_fan_h[k],fan_h_dict[k]))
        erros = 0
        es = eq
        es.extend(edh)
        es.extend(sth)
        es.extend(fanh)
        # 求各相对误差的平方误差,2020/3/11测试结果比较好
        erros = math.sqrt(sum([er**2 for er in es]))
        return erros

    def fitness(self, X):
        X = self.decode(X)
        f = self.rela_erros(X)
        return f
    
    def generator(self, random, args):
        return self.rand(random)
        
    def evaluator(self, candidates, args):
        fitness = []
        for X in candidates:
            score = self.fitness(X)
            #添加到列表
            fitness.append(score)
        return fitness

def _my_avg_terminator(population, num_generations, num_evaluations, args):
    if num_generations > args['max_generations']:
        return True
    tolerance = args.setdefault('tolerance', 0.001)
    pop_fit = [x.fitness for x in population]
    avg_fit = sum(pop_fit) / float(len(population))
    best_fit = max(pop_fit)
    clear_dis()
    msg = '%.1lf%%'%((num_generations)*1.0/(args['max_generations'])*100)
    print('正在计算...完成%s'%msg)
    print('计算次数[第%d代]:%d'%(num_generations+1,num_evaluations))
    if abs(best_fit - avg_fit) < tolerance and best_fit < args['minimum_fitness']:
        print(best_fit, avg_fit, best_fit)
        return True
    return False
    
def _my_terminator(population, num_generations, num_evaluations, args):
    if num_generations > args['max_generations']:
        return True
    for x in population:
        if x.fitness < 1e-5:
            return True
    return False

# 自定义终止条件
# 根据指定的类型来选择终止条件函数
def my_terminator(terminator_type):
    TerminatorDict = {
        # 按进化代数终止: 表示进化多少代之后就终止,无论是否达到最优,也不考虑算法是否收敛
        1:terminators.generation_termination,
        # 按适应值计算次数终止: 表示适应值计算多少次后就就终止,无论是否达到最优,也不考虑算法是否收敛
        2:terminators.evaluation_termination,
        # 按种群的平均适应值终止: 表示种群的平均适应值连续多少代没有显著变化就终止,因为进化过程已经收敛或者假收敛了
        3:terminators.average_fitness_termination,
        4:_my_terminator,
        5:terminators.no_improvement_termination,
        6:_my_avg_terminator
    }
    if not (terminator_type in TerminatorDict):
        terminator_type = 1
    return TerminatorDict[terminator_type]

def generate_rastrigin(random, args):
    return args['problem'].generator(random, args)

def evaluate_rastrigin(candidates, args):
    return args['problem'].evaluator(candidates, args)

# 自定义个体生成器
# 注: 下面es.evolve()函数中的所有参数构成了一个词典args
def my_generator(random, args):
    # 跳转到Problem::rand()成员函数
    return args['problem'].rand(random)

# 自定义个体适应值计算函数
def my_fitness(X, args):
    # 跳转到Problem::fitness()成员函数
    return args['problem'].fitness(X)

# 自定义种群适应值计算函数
# 注: 参数candidates是一个数组,它由多个个体组成
def my_evaluator(candidates, args):
    #依次计算每个个体的适应值
    fitness = []
    for X in candidates:
        # 将浮点数四舍五入,强制转换为0和1,并计算个体的适应值
        score = my_fitness(X, args)
        #添加到列表
        fitness.append(score)
    return fitness

# 遗传算法主框架
def runGA(args):
    #随机数初始化
    rng = random.Random()
    rng.seed(int(time.time()))
    algo = args.get('algo', 1)
    es = None
    if algo == 'ga':
        # 遗传算法
        es = ec.GA(rng)
    elif algo == 'pso':
        # 粒子群算法
        es = swarm.PSO(rng)
        es.topology = swarm.topologies.ring_topology
    elif algo == 'sa':
        # 模拟退火算法
        es = ec.SA(rng)
    elif algo == 'dea':
        # 差分进化算法
        es = ec.DEA(rng)
    elif algo == 'es':
        # 进化策略算法
        es = ec.ES(rng)    
    else:
        # 默认采用遗传算法
        es = ec.GA(rng)

    #设置终止条件
    es.terminator = my_terminator(args['terminator'])
    if args['use_observe_file']:
        es.observer = observers.file_observer
        # 新建种群数据记录文件
        stat_file = open(args['stat_file'], 'w')
        ind_file = open(args['ind_file'], 'w')
    else:
        stat_file = None
        ind_file = None
    problem = Problem(args)
    if args['use_palla']:
        final_pop = es.evolve(
                            generator=generate_rastrigin, 
                            evaluator=inspyred.ec.evaluators.parallel_evaluation_mp,
                            mp_evaluator=evaluate_rastrigin, 
                            mp_nprocs=args.get('mp_nprocs',4),
                            pop_size=args.get('popsize', 10), 
                            bounder=inspyred.ec.Bounder(args['lb'],args['ub']),
                            maximize=False,
                            max_evaluations = args.get('max_evaluations', 1200),   # 最大适应值计算次数
                            max_generations = args.get('max_generations', 200),
                            statistics_file = stat_file, 
                            individuals_file = ind_file, 
                            problem = problem,
                            # 遗传算法GA参数
                            mutation_rate = args.get('mutation_rate',0.015),       # 变异率
                            cross_rate =  args.get('cross_rate',0.75),            # 交叉率
                            # 粒子群算法PSO参数
                            neighborhood_size =  args.get('neighborhood_size',5),  # 邻域范围
                            minimum_fitness = args.get('minimum_fitness',0.1),
                            #   差分进化
                            crossover_rate = args.get('cross_rate',0.75))
    
    else:
        final_pop = es.evolve(
                            generator = my_generator,                              # 自定义个体生成函数
                            evaluator = my_evaluator,                              # 自定义个体适应值计算函数
                            pop_size=args.get('popsize', 10), 
                            bounder=inspyred.ec.Bounder(args['lb'],args['ub']),
                            maximize=False,
                            max_evaluations = args.get('max_evaluations', 1200),   # 最大适应值计算次数
                            max_generations = args.get('max_generations', 20),
                            statistics_file = stat_file, 
                            individuals_file = ind_file, 
                            problem = problem,
                            # 遗传算法GA参数
                            mutation_rate = args.get('mutation_rate',0.015),       # 变异率
                            cross_rate =  args.get('cross_rate',0.75),            # 交叉率
                            # 粒子群算法PSO参数
                            neighborhood_size =  args.get('neighborhood_size',5),  # 邻域范围
                            minimum_fitness = args.get('minimum_fitness',0.1),
                            #   差分进化
                            crossover_rate = args.get('cross_rate',0.75))

    if args['use_observe_file']:
        # 关闭种群数据记录文件
        stat_file.close()
        ind_file.close()
    final_pop.sort(reverse = True)
    return final_pop[0].candidate

def init_config(vnet,algo,init_args,max_generations=10,popsize=50,cross_rate=0.9,mutation_rate=0.5):
    
    # 参数设置
    config = {
        #---------------- 自定义参数: 变量及范围  -----------------#
        # 变量类型(0-整数,1-实数) 
        "var_types":init_args['vt'],    
        # 下限(数组中的每一位代表一个变量的下限)
        "lb":init_args['lbs'],
        # 上限(数组中的每一位代表一个变量的上限)
        "ub":init_args['ubs'],
        "vnet":vnet,
        "Q0":init_args['init_Q0'],
        "H0":init_args['init_H0'],
        "P0":init_args['init_P0'],
        'sel_params':init_args['sel_params'],
        "edge_ids":init_args['ids'],
        'struct_eids':init_args['struct_eids'],
        "weight":init_args['weight'],
        #---------------- 自定义参数: 距离度量方法及终止条件  -----------------#
        "algo": algo,                      # 算法选择
                                        #   1. 遗传算法GA
                                        #   2. 粒子群算法PSO
                                        #   3. 模拟退火算法SA
                                        #   4. 差分进化算法DEA
                                        #   5. 进化策略ES
                
        "terminator": 1,                # 终止条件(1-表示按最大进化代数终止迭代)
                                        #   1. 按最大进化代数终止迭代
                                        #   2. 按最大适应值计算次数终止迭代
                                        #   3. 按种群平均适应值收敛终止迭代

        #---------------- 进化算法通用参数(遗传、粒子群、蚁群、模拟退火等)  -----------------#
        "popsize": popsize,                  # 种群规模
        "max_generations": max_generations,        # 最大进化代数
        "max_evaluations": 1000,        # 最大适应值计算次数(暂未使用)

        #---------------- 遗传算法GA参数 --------------------#
        "cross_rate": cross_rate,             # 交叉率(默认1.0)
        "mutation_rate": mutation_rate,         # 变异率(默认0.1)

        #---------------- 粒子群算法PSO参数 -----------------#
        "neighborhood_size": 5,          # 邻域范围
    }
    return config

def ea_test(config):
    start = time.time()
    # 遗传算法计算得到"最优"个体
    best_X = runGA(config)
    # 结束时间
    stop = time.time()
    print('计算耗时:', "%.1fs" % (stop-start))
    return best_X

def ea_sensor(pro_datas, graph_datas,args):
    out_dir = args['result_dir']
    if args['use_observe_file']:
        creat_dir(out_dir)
    algo = args['algo'] 
    #构造通风网络,读取并设置相关数据
    vnet = VentNetwork()
    if not build_network(graph_datas, vnet.graph()):
        print('\t-->构造通风网络失败!!!')
        return False
    #添加虚拟源汇,将网络变成单一源汇通风网络
    vnet.addVirtualST()

    dg = vnet.graph()
    edges = dg.es
    edge_ids = []
    lbs = []
    ubs = []
    struct_eids = []
    edge_set = pro_datas['edge_set']
    node_set = pro_datas['node_set']
    struct_set = pro_datas['struct_set']
    fan_set = pro_datas['fan_set']
    weight = {}
    for ee in edges:
        if is_zero_edge(ee):
            continue
        eid = ee['id']
        edge_ids.append(eid)
        pro_e = edge_set[eid]
        lbs.append(pro_e['r_lb'])
        ubs.append(pro_e['r_ub'])
        weight[pro_e['id']]=1
    # 构筑物
    for k,st in struct_set.items():
        if st['r_lb'] is not None and st['r_ub'] is not None:
            struct_eids.append(st['e_id'])
            lbs.append(st['r_lb'])
            ubs.append(st['r_ub'])
    vt = [1]*(len(edge_ids)+len(struct_eids))
    # 初始化参数
    init_args = {
    'vt':vt,
    'ids':edge_ids,
    'struct_eids':struct_eids,
    'lbs':lbs,
    'ubs':ubs,
    'init_Q0':None,
    'init_H0':None,
    'init_P0':None,
    'weight':weight,
    'sel_params':args['sel_params']
    }
    config = init_config(vnet,algo,init_args,max_generations=args['max_generations'],popsize=args['popsize'],cross_rate=args['cross_rate'],mutation_rate=args['mutation_rate'])
    config['stat_file'] = out_dir+"stat.csv"
    config['ind_file'] = out_dir+"ind.csv"
    config['use_observe_file'] = args['use_observe_file']
    config['pro_datas'] = pro_datas
    config['terminator'] = args['terminator']
    config['minimum_fitness'] = args['minimum_fitness']
    config['use_palla'] = args['use_palla']
    config['mp_nprocs'] = args['mp_nprocs']
    print('开始计算....')
    best_X = ea_test(config)
    best_X = [round(x,DOUBLE_PRECISION) for x in best_X]
    if args['use_observe_file']:
        ea_utils.plot(config['stat_file'], out_dir+'ea.png', False)
    print("数据写入中...")
    best_edge_rs = best_X[0:len(edge_ids)]
    best_struct_rs = best_X[len(edge_ids):len(edge_ids)+len(struct_eids)]
    print(u"分支最优解：",best_edge_rs)
    print(u"构筑物最优解：",best_struct_rs)
    q_dict,delta_h_dict,st_h_dict, fan_h_dict = each_vno(vnet,edge_ids,struct_eids,best_edge_rs,best_struct_rs,sel_params=args['sel_params'], is_print = True,all_q=True)
    # print_network(vnet, msg='网络解算结果')
    if (q_dict is None) and (delta_h_dict is None):
        vnet.delVirtualST()
        return False
    for index,eid in enumerate(edge_ids):
        e = edge_set[eid]
        e['q']=round(q_dict[eid],DOUBLE_PRECISION)
        e['delta_r']=round(best_edge_rs[index],DOUBLE_PRECISION)
        e['r']=round(e['delta_r']+e['r0'],DOUBLE_PRECISION)

    for index,st_id in enumerate(struct_set):
        st = struct_set[st_id]
        if index >= len(best_struct_rs):
            continue
        st['delta_r']=round(best_struct_rs[index],DOUBLE_PRECISION)
        if st['r0'] is not None:
            st['r'] = round(st['delta_r']+st['r0'],DOUBLE_PRECISION)
    
    for index,eid in enumerate(edge_ids):
        e = edge_set[eid]
        st_delta_r = 0
        st_r0 = 0
        for index,st_id in enumerate(struct_set):
            st = struct_set[st_id]
            if st['e_id'] is not None and st['e_id'] == eid and st['r0'] is not None and st['delta_r'] is not None:
                st_delta_r = st['delta_r']
                st['q'] = e['q']
                st_r0 = st['r0']
                # r = st_delta_r+e['r']+st_r0
                r = st_delta_r+st_r0
                h = r*e['q']*abs(e['q'])
                st['h'] = round(h,DOUBLE_PRECISION)
                break
        # 分支的风阻等于构筑物的初始风阻+构筑物变化风阻+分支初始风阻+分支变化风阻
        e['r'] = round(st_delta_r+e['r']+st_r0,DOUBLE_PRECISION)
        e['h'] = round(e['r']*e['q']*abs(e['q']),DOUBLE_PRECISION)

    nodes = dg.vs
    for vv in nodes:
        if vnet.sn != vv.index  and  vnet.tn != vv.index:
            pro_v = node_set[vv['id']]
            pro_v['p']=round(vv['p'],DOUBLE_PRECISION)
    # 属性数据的格式{'id':1,'id_v_high':7,'id_v_low':3,'sensor_delta_h':111,''}
    # delta_h_dict的格式{'高压节点id-低压节点id':传感器(或给定)的压差}
    pro_delta_h = pro_datas['sensor_hs_set']
    for k,value in pro_delta_h.items():
        dh_dict = pro_delta_h[k]
        lu_ids = '%s-%s'%(dh_dict['s'],dh_dict['t'])
        if lu_ids in delta_h_dict:
            dh_dict['delta_h'] = round(delta_h_dict[lu_ids],DOUBLE_PRECISION)

        if dh_dict['idFan'] is not None:
            # 风机压差
            for _,fan in fan_set.items():
                if dh_dict['idFan'] == fan['id']:
                    # 格式：{'末节点id-分支id':压差}
                    lu_ids = '%s-%s'%(dh_dict['t'],fan['e_id'])
                    if lu_ids in fan_h_dict:
                        dh_dict['delta_h'] = round(fan_h_dict[lu_ids],DOUBLE_PRECISION)

        if dh_dict['idSt'] is not None:
            # 构筑物压差
            for _,st in struct_set.items():
                if dh_dict['idSt'] == st['id']:
                    if st['e_id'] in st_h_dict:
                        dh_dict['delta_h'] = round(st_h_dict[st['e_id']],DOUBLE_PRECISION)

    pro_delta_q = pro_datas['sensor_qs_set']
    for k,value in pro_delta_q.items():
        dq_dict = value
        if value['e_id'] in q_dict:
            dq_dict['cacul_q'] = round(q_dict[value['e_id']],DOUBLE_PRECISION)

    # 删除虚拟源汇
    vnet.delVirtualST()
    return True
